﻿using System;
using System.Collections.Generic;
using Algoritmia.DataStructures.Maps;
using Algoritmia.SemiRings;
using Algoritmia.StateSpaces;
using Algoritmia.Utils;

namespace Algoritmia.Schemes.DynamicProgramming
{
    public class IterativeIdempotentDynamicProgrammingSolver<TState,TDecision,TResult,TWeight> : 
        IterativeDynamicProgrammingSolver<TState,TDecision,TResult,TWeight>
    {
        private readonly EqualityComparison<TResult> _equalityComparison;

        public IterativeIdempotentDynamicProgrammingSolver(Func<IBackwardsStateSpace<TState, TDecision>, IMap<TState, TResult>> createMap = null,
                                                           EqualityComparison<TResult> equalityComparison = null)
            : base(createMap)
        {
            _equalityComparison = AlgoritmiaFactory.GetDefaultEqualityComparison(equalityComparison);
        }

        public IEnumerable<KeyValuePair<TState,TDecision>> Backpointers(IBackwardsStateSpace<TState,TDecision> space,
            IIdempotentSemiRing<TResult,TWeight> semiRing, Func<TState,TDecision,TWeight> wfn)
        {
            var mem = _createMap(space);
            var topsorter = _createStateSpaceTopsorter();
            foreach(var s in topsorter.Topsorted(space))
            {
                TDecision decision = default(TDecision);
                if ( space.IsInitial(s))
                {
                    mem[s] = semiRing.One;
                }
                else
                {
                    mem[s] = semiRing.Zero;
                }
                foreach(var d in space.IncomingDecisions(s))
                {
                    var s2 = space.Undo(s, d);
                    var score = semiRing.Times(mem[s2], wfn(s2,d));
                    if ( !_equalityComparison(semiRing.Plus(mem[s],score),mem[s]) )
                    {
                        mem[s] = score;
                        decision = d;
                    }
                }
                yield return new KeyValuePair<TState, TDecision>(s,decision);
            }
        }
    }

    public class IterativeIdempotentDynamicProgrammingSolver<TState,TDecision,TResult> : 
        IterativeIdempotentDynamicProgrammingSolver<TState,TDecision,TResult,TResult>
    {
        public IterativeIdempotentDynamicProgrammingSolver(Func<IBackwardsStateSpace<TState, TDecision>, IMap<TState, TResult>> createMap = null,
                                                           EqualityComparison<TResult> equalityComparison = null)
            : base(createMap,equalityComparison)
        {
        }
    }
}
